import numpy as np
import tensorflow as tf
from tensorflow.keras.layers import Dropout, Dense, SimpleRNN
import matplotlib.pyplot as plt
import pandas as pd
from sklearn.preprocessing import MinMaxScaler
from sklearn.metrics import mean_squared_error, mean_absolute_error
import math
import os
import shutil

# 获取当前脚本所在的绝对路径
script_dir = os.path.dirname(os.path.abspath(__file__))
print(f"脚本目录: {script_dir}")

# 强制使用CPU
os.environ['CUDA_VISIBLE_DEVICES'] = '-1'

print(f"NumPy version: {np.__version__}")
print(f"TensorFlow version: {tf.__version__}")
print(f"GPU available: {tf.config.list_physical_devices('GPU')}")

# 使用正确的NumPy对象检查
print("np.object_ check:", hasattr(np, 'object_'))

# 加载数据 - 使用绝对路径
data_path = os.path.join(script_dir, 'SH600519.csv')
print(f"数据路径: {data_path}")
maotai = pd.read_csv(data_path)
training_set = maotai.iloc[0:2426 - 300, 2:3].values
test_set = maotai.iloc[2426 - 300:, 2:3].values

# 归一化
sc = MinMaxScaler(feature_range=(0, 1))
training_set_scaled = sc.fit_transform(training_set)
test_set = sc.transform(test_set)

# 构建训练数据
x_train, y_train = [], []
for i in range(60, len(training_set_scaled)):
    x_train.append(training_set_scaled[i - 60:i, 0])
    y_train.append(training_set_scaled[i, 0])

# 打乱数据
np.random.seed(7)
indices = np.arange(len(x_train))
np.random.shuffle(indices)
x_train = np.array(x_train)[indices]
y_train = np.array(y_train)[indices]
x_train = np.reshape(x_train, (x_train.shape[0], 60, 1))

# 构建测试数据
x_test, y_test = [], []
for i in range(60, len(test_set)):
    x_test.append(test_set[i - 60:i, 0])
    y_test.append(test_set[i, 0])
x_test, y_test = np.array(x_test), np.array(y_test)
x_test = np.reshape(x_test, (x_test.shape[0], 60, 1))

# 模型定义
model = tf.keras.Sequential([
    SimpleRNN(80, return_sequences=True),
    Dropout(0.2),
    SimpleRNN(100),
    Dropout(0.2),
    Dense(1)
])

model.compile(optimizer=tf.keras.optimizers.Adam(0.001),
              loss='mean_squared_error')

# 解决方案1：直接使用脚本目录而不创建子目录
checkpoint_save_path = os.path.join(script_dir, "rnn_stock.ckpt")
print(f"使用脚本目录作为检查点路径: {checkpoint_save_path}")

# 解决方案2：使用临时目录
# import tempfile
# checkpoint_dir = tempfile.mkdtemp()
# checkpoint_save_path = os.path.join(checkpoint_dir, "rnn_stock.ckpt")
# print(f"使用临时目录作为检查点路径: {checkpoint_save_path}")

# 确保文件目录存在
checkpoint_dir = os.path.dirname(checkpoint_save_path)
if not os.path.exists(checkpoint_dir):
    os.makedirs(checkpoint_dir, exist_ok=True)

# 检查点回调 - 使用绝对路径
cp_callback = tf.keras.callbacks.ModelCheckpoint(
    filepath=checkpoint_save_path,
    save_weights_only=True,
    save_best_only=True,
    monitor='val_loss'
)

# 如果存在检查点文件，加载权重
if os.path.exists(checkpoint_save_path + '.index'):
    print('-------------加载模型-----------------')
    model.load_weights(checkpoint_save_path)

# 模型训练
print("开始训练模型...")
try:
    history = model.fit(
        x_train, y_train,
        batch_size=64,
        epochs=50,
        validation_data=(x_test, y_test),
        validation_freq=1,
        callbacks=[cp_callback]
    )
except Exception as e:
    print(f"训练过程中发生错误: {e}")
    print("尝试不使用回调函数训练...")
    history = model.fit(
        x_train, y_train,
        batch_size=64,
        epochs=50,
        validation_data=(x_test, y_test),
        validation_freq=1
    )

# 后续代码保持不变...
model.summary()

# 保存权重 - 使用绝对路径
weights_file = os.path.join(script_dir, 'weights.txt')
with open(weights_file, 'w') as file:
    for v in model.trainable_variables:
        file.write(f"{v.name}\n{v.shape}\n{v.numpy()}\n\n")
print(f"权重已保存到: {weights_file}")
# 方案1: 使用系统字体
plt.rcParams['font.sans-serif'] = ['SimHei']  # 设置黑体
plt.rcParams['axes.unicode_minus'] = False  # 正确显示负号
print("使用系统黑体显示中文")
# 绘制损失曲线
loss = history.history['loss']
val_loss = history.history['val_loss']
plt.figure(figsize=(10, 6))
plt.plot(loss, label='训练损失')
plt.plot(val_loss, label='验证损失')
plt.title('训练和验证损失')
plt.xlabel('迭代次数')
plt.ylabel('损失值')
plt.legend()

# 保存图像 - 使用绝对路径
loss_image = os.path.join(script_dir, 'loss_curve.png')
plt.savefig(loss_image)
print(f"损失曲线图已保存到: {loss_image}")
plt.show()

# 预测和评估
predicted_stock_price = model.predict(x_test)
predicted_stock_price = sc.inverse_transform(predicted_stock_price)
real_stock_price = sc.inverse_transform(test_set[60:])

plt.figure(figsize=(12, 6))
plt.plot(real_stock_price, color='red', label='茅台股票价格')
plt.plot(predicted_stock_price, color='blue', label='预测的茅台股票价格')
plt.title('茅台股票价格预测')
plt.xlabel('时间')
plt.ylabel('价格')
plt.legend()

# 保存图像 - 使用绝对路径
prediction_image = os.path.join(script_dir, 'stock_prediction.png')
plt.savefig(prediction_image)
print(f"股票预测图已保存到: {prediction_image}")
plt.show()

# 计算评估指标
mse = mean_squared_error(real_stock_price, predicted_stock_price)
rmse = math.sqrt(mse)
mae = mean_absolute_error(real_stock_price, predicted_stock_price)
print(f'均方误差: {mse:.6f}')
print(f'均方根误差: {rmse:.6f}')
print(f'平均绝对误差: {mae:.6f}')

# 如果有临时目录，清理它
if 'tempfile' in locals() and os.path.exists(checkpoint_dir):
    print(f"清理临时目录: {checkpoint_dir}")
    shutil.rmtree(checkpoint_dir)